package com.demo.util;

import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import javax.sql.DataSource;
import java.sql.Connection;
import java.sql.SQLException;

/**
 * @ClassName ConnectionUtils
 * @Description TODO
 * @Author Windy
 * @Date 2022/8/31 15:02
 * @Version 1.0
 **/
@Component
public class ConnectionUtils {

    private ThreadLocal<Connection> threadLocal= new ThreadLocal<>();
    DataSource dataSource;

    public void setDataSource(DataSource dataSource) {
        this.dataSource = dataSource;
    }

    public Connection getConnection() {
        Connection conn = threadLocal.get();
        if (conn == null) {
            try {
                conn = dataSource.getConnection();
                threadLocal.set(conn);
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
        return conn;
    }

    public void removeConnection() {
        threadLocal.remove();
    }
}
